﻿using Unity.Mathematics;
using UnityEngine;
using UnityEngine.Profiling;

namespace Mars.GpuInstance
{
    public class GPUAnimator
    {

        public GPUAnimationClipData[] currentAnimationClipData;
        public Vector2 currentAnimationClipDataWeights;
        public int activeClipCount;
        
        public GPUAnimationClipData[] newAnimationClipData;

        public Vector2 currentClipIndices;
        public Vector2 currentClipStartTimes;
        public Vector2 newClipStartTimes;
        public Vector2 currentClipSpeeds;
        
        public int transitionIndex;
        public GPUTransition transition;

        public bool nonPropScaling;

        // private AnchorInfo[] anchorInfos;
        //
        // private Dictionary<int, int> anchorNameToBoneIndex;
        
        public GPUAnimator()
        {
            ResetAnimator();
        }
        
        public void ResetAnimator()
        {
            if (currentAnimationClipData == null)
            {
                currentAnimationClipData = new GPUAnimationClipData[2];
                newAnimationClipData = new GPUAnimationClipData[2];
            }
            else
            {
                for (int i = 0; i < 2; i++)
                {
                    currentAnimationClipData[i] = default;
                    newAnimationClipData[i] = default;
                }

                transition = null;
            }

            currentClipIndices = Vector2.zero;
            currentClipSpeeds = Vector2.one;
            currentAnimationClipDataWeights = Vector2.zero;
            currentClipStartTimes = Vector2.zero;
            newClipStartTimes = Vector2.zero;
            transitionIndex = -1;
        }

        //更新挂点信息
        // public void UpdateAnchorInfos(AnchorInfo[] _anchorInfos)
        // {
        //     anchorInfos = _anchorInfos;
        // }
        //
        // public void UpdateAnchorNameToBoneIndex(Dictionary<int, int> _anchorNameToBoneIndex)
        // {
        //     anchorNameToBoneIndex = _anchorNameToBoneIndex;
        // }
        
        public void StartAnimation(GPUAnimatorInstancerRuntimeData runtimeData, int arrayIndex, AnimationClip animationClip, float startTime = -1.0f, float speed = 1.0f)
        {
            if (runtimeData.animationClipDataDict.TryGetValue(animationClip.GetHashCode(), out GPUAnimationClipData clipData))
            {
                StartAnimation(runtimeData, arrayIndex, clipData, startTime, speed);
            }
            else
            {
                Debug.LogError($"没有找到{animationClip.name}对应的烘焙数据，重新烘焙试试:");
            }
        }

        public void StartAnimation(GPUAnimatorInstancerRuntimeData runtimeData, int arrayIndex, GPUAnimationClipData clipData, float startTime = -1.0f, float speed = 1.0f)
        {
            Profiler.BeginSample("StartAnimation 1");
            float currentTime = Time.time;
            
            int previousClipIndex = GetCurrentClipIndex(clipData);

            if (speed <= 0)
            {
                speed = 0.000001f;
            }

            if (startTime >= 0)
            {
                newClipStartTimes.x = currentTime - startTime;
            }
            else if (previousClipIndex >= 0)
            {
                if (speed == currentClipSpeeds[previousClipIndex])
                {
                    newClipStartTimes.x = currentClipStartTimes[previousClipIndex];
                }
                else
                {
                    newClipStartTimes.x = GetSpeedRelativeStartTime(currentTime, currentClipStartTimes[previousClipIndex], currentClipSpeeds[previousClipIndex], clipData.length, speed);
                }
            }
            else
            {
                newClipStartTimes.x = currentTime;
            }
            newClipStartTimes.y = 0;
            //newClipStartTimes.z = 0;
            //newClipStartTimes.w = 0;
            
            currentClipIndices.y = 0;
            currentClipIndices.x = clipData.clipIndex;
            
            currentAnimationClipData[1] = default;
            currentAnimationClipData[0] = clipData;
            
            currentAnimationClipDataWeights.y = 0;
            currentAnimationClipDataWeights.x = 1;
            currentClipSpeeds[0] = speed;
            activeClipCount = 1;

            nonPropScaling = clipData.nonPropScaling;
            
            Profiler.EndSample();
            
            Profiler.BeginSample("StartAnimation 2");
            BlendAnimations(runtimeData, arrayIndex);
            Profiler.EndSample();
        }

        public void StartCrossFade(GPUAnimatorInstancerRuntimeData runtimeData, int arrayIndex, GPUAnimationClipData clipData, float startTime, float transitionTime)
        {
            if (activeClipCount == 0)
            {
                StartAnimation(runtimeData, arrayIndex, clipData);
                return;
            }

            var prevClipData = currentAnimationClipData[0];
            var prevDataIndex = 0;
            if (activeClipCount > 1)
            {
                var weights = currentAnimationClipDataWeights;
                if (transitionIndex >= 0 && transition != null)
                    weights = transition.endWeights;
                var maxWeight = -1f;
                var maxIndex = -1;
                for (var i = 0; i < 2; ++i)
                {
                    if (weights[i] > maxWeight)
                    {
                        maxIndex = i;
                        maxWeight = weights[i];
                    }
                }

                if (maxIndex > 0)
                {
                    prevClipData = currentAnimationClipData[maxIndex];
                    prevDataIndex = maxIndex;
                }
            }

            if (prevClipData.clipIndex == clipData.clipIndex)
            {
                return;
            }

            // 永远把最新的 clip 放在 0 位置。方便 shader 内根据机型去除 CrossFade 效果
            currentClipIndices.x = clipData.clipIndex;
            currentClipIndices.y = prevClipData.clipIndex;

            var animationTimes = new float2(
                startTime,
                (Time.time - currentClipStartTimes[prevDataIndex]) * currentClipSpeeds[prevDataIndex]
            );

            var animationSpeeds = new float2(
                currentClipSpeeds[prevDataIndex], 
                currentClipSpeeds[prevDataIndex]
            );
            nonPropScaling = clipData.nonPropScaling;

            StartBlend(runtimeData, arrayIndex, prevClipData, clipData, animationTimes, animationSpeeds, transitionTime);
        }

        public void StartBlend(GPUAnimatorInstancerRuntimeData runtimeData, int arrayIndex, GPUAnimationClipData fromClip, GPUAnimationClipData toClip,  float2 animationTimes, float2 animationSpeeds, float transitionTime = 0)
        {
            float currentTime = Time.time;
            bool hasTransition = transitionTime > 0;
            int previousClipCount = activeClipCount;

            activeClipCount = 1;

            if (hasTransition && previousClipCount > 0)
            {
                activeClipCount = 2;
            }
            
            newAnimationClipData[0] = toClip;
            newAnimationClipData[1] = fromClip;
            
            for (int i = 0; i < 2; i++)
            {
                newClipStartTimes[i] = animationTimes[i];
            }

            for (int i = 0; i < 2; i++)
            {
                currentAnimationClipData[i] = newAnimationClipData[i];
            }
            
            Vector2 newWeights = new Vector2(1.0f,0);
            if (hasTransition)
            {
                currentAnimationClipDataWeights = new Vector2(0, 1.0f);

                if (transition == null)
                {
                    transition = new GPUTransition();
                }
                transition.SetData(arrayIndex, currentTime, transitionTime, previousClipCount, currentAnimationClipDataWeights, newWeights,  activeClipCount - previousClipCount);
                if (transitionIndex < 0)
                {
                    transitionIndex = runtimeData.transitioningAnimators.Count;
                    runtimeData.transitioningAnimators.Add(this);
                }
            }
            else
            {
                currentAnimationClipDataWeights = newWeights;
            }

            int speedCount = 2;

            for (int i = 0; i < 2; i++)
            {
                if (animationSpeeds[i] <= 0)
                {
                    currentClipSpeeds[i] = 0.000001f;
                }
                else
                {
                    currentClipSpeeds[i] = animationSpeeds[i];
                }

                newClipStartTimes[i] = currentTime - (newClipStartTimes[i] / currentClipSpeeds[i]);
            }

            BlendAnimations(runtimeData, arrayIndex, hasTransition);
        }
        
        public bool ApplyTransition(GPUAnimatorInstancerRuntimeData runtimeData, float currentTime)
        {
            if (transitionIndex < 0 || transition == null)
            {
                return false;
            }

            int animationDataIndex = transition.arrayIndex * 2;
            if (animationDataIndex >= runtimeData.animationData.Length)
            {
                return false;
            }

            if (currentTime > transition.startTime + transition.totalTime)
            {
                activeClipCount = transition.endActiveClipCount;
                currentAnimationClipDataWeights = transition.endWeights;
                RemoveFromTransitioningAnimatorsSwapBack(runtimeData);
            }
            else
            {
                currentAnimationClipDataWeights = Vector4.Lerp(transition.startWeights, transition.endWeights, (currentTime - transition.startTime) / transition.totalTime);
            }

            // set weights
            var animationData = runtimeData.animationData[animationDataIndex];
            animationData.y = currentAnimationClipDataWeights.x;
            animationData.w = currentAnimationClipDataWeights.y;
            runtimeData.animationData[animationDataIndex] = animationData;
            return transitionIndex >= 0;
        }
        
        private void BlendAnimations(GPUAnimatorInstancerRuntimeData runtimeData, int arrayIndex, bool hasTransition = false)
        {
            currentClipStartTimes = newClipStartTimes;
            newClipStartTimes = Vector4.zero;
            
            ApplyToRuntimeData(runtimeData, arrayIndex, hasTransition);
        }

        internal void ApplyToRuntimeData(GPUAnimatorInstancerRuntimeData runtimeData, int arrayIndex, bool hasTransition = false)
        {
            if (activeClipCount <= 0)
            {
                return;
            }
            
            if (!hasTransition && transitionIndex >= 0)
            {
                RemoveFromTransitioningAnimatorsSwapBack(runtimeData);
            }

            int animationIndex = arrayIndex * 2;

            GPUAnimationClipData clipData;
            float2 animationFrames = float2.zero;
            for (int i = 0; i < 2; i++)
            {
                if (i < activeClipCount && currentAnimationClipData[i].length > 0)
                {
                    clipData = currentAnimationClipData[i];

                    // x ->  minFrame, y -> maxFrame, z -> speed, w -> startTime
                    runtimeData.animatorControllerData[animationIndex + i] = new float4(
                        clipData.clipStartFrame,
                        // 非循环动画为负数
                        clipData.IsLoopDisabled() ? -(clipData.clipStartFrame + clipData.clipFrameCount - 2) : clipData.clipStartFrame + clipData.clipFrameCount - 1,
                        currentClipSpeeds[i],
                        currentClipStartTimes[i]);

                    // set current clip frame
                    var clipFrame = GetClipFrame(i, Time.time, clipData.length, clipData.clipFrameCount, clipData.IsLoopDisabled());
                    animationFrames[i] = clipFrame;
                    //animationFrames[i * 2] = clipData.clipStartFrame + clipFrame;
                    //animationFrames[i * 2 + 1] = clipData.clipNonPropScaleStartFrame + clipFrame;
                }
                else
                {
                    runtimeData.animatorControllerData[animationIndex + i] = float4.zero;
                    animationFrames[i] = -1.0f;
                    // animationFrames[i * 2] = -1.0f;
                    // animationFrames[i * 2 + 1] = -1.0f;
                }
            }
            // frame, weight tuple
            runtimeData.animationData[animationIndex] = new float4(animationFrames.x, currentAnimationClipDataWeights.x, animationFrames.y, currentAnimationClipDataWeights.y);
            // set clip indices
            runtimeData.animationData[animationIndex + 1] = new float4(currentClipIndices.x, currentClipIndices.y, 0, 0);

            runtimeData.animationDataModified = true;
            runtimeData.animatorDataModified = true;

            if (runtimeData.inputNonPropScaling[arrayIndex] != nonPropScaling)
            {
                runtimeData.inputNonPropScaling[arrayIndex] = nonPropScaling;
                runtimeData.dataModified = true;
            }
        }

        public void SetAnimationSpeed(GPUAnimatorInstancerRuntimeData runtimeData, int arrayIndex, float animationSpeed)
        {
            int crowdAnimIndex = arrayIndex * 2;
            for (int i = 0; i < activeClipCount; i++)
            {
                Vector4 animatorData = runtimeData.animatorControllerData[crowdAnimIndex + i];
                if (animationSpeed <= 0)
                {
                    animationSpeed = 0.000001f;
                }
                currentClipStartTimes[i] = GetSpeedRelativeStartTime(Time.time, currentClipStartTimes[i], currentClipSpeeds[i], currentAnimationClipData[i].length, animationSpeed);
                animatorData.w = currentClipStartTimes[i];
                currentClipSpeeds[i] = animationSpeed;
                animatorData.z = animationSpeed;

                runtimeData.animatorControllerData[crowdAnimIndex + i] = animatorData;
            }
            runtimeData.animatorDataModified = true;
        }
        
        public void SetAnimationSpeeds(GPUAnimatorInstancerRuntimeData runtimeData, int arrayIndex, float[] animationSpeeds)
        {
            int speedCount = 0;
            if (animationSpeeds != null)
            {
                speedCount = animationSpeeds.Length;
            }
            int crowdAnimIndex = arrayIndex * 2;
            for (int i = 0; i < activeClipCount; i++)
            {
                Vector4 animatorData = runtimeData.animatorControllerData[crowdAnimIndex + i];
                if (animationSpeeds[i] <= 0)
                {
                    animationSpeeds[i] = 0.000001f;
                }
                currentClipStartTimes[i] = GetSpeedRelativeStartTime(Time.time, currentClipStartTimes[i], currentClipSpeeds[i], currentAnimationClipData[i].length, animationSpeeds[i]);
                animatorData.w = currentClipStartTimes[i];
                if (i < speedCount)
                {
                    currentClipSpeeds[i] = animationSpeeds[i];
                }
                else
                {
                    currentClipSpeeds[i] = 1;
                }
                animatorData.z = currentClipSpeeds[i];

                runtimeData.animatorControllerData[crowdAnimIndex + i] = animatorData;
            }
            runtimeData.animatorDataModified = true;
        }
        
        public void SetAnimationWeights(GPUAnimatorInstancerRuntimeData runtimeData, int arrayIndex, Vector2 animationWeights)
        {
            int animationDataIndex = arrayIndex * 2;

            // set weights
            currentAnimationClipDataWeights = animationWeights;
            var animationData = runtimeData.animationData[animationDataIndex];
            animationData.y = currentAnimationClipDataWeights.x;
            animationData.w = currentAnimationClipDataWeights.y;
            runtimeData.animationData[animationDataIndex] = animationData;
            runtimeData.animationDataModified = true;
        }
        
        public GPUAnimationClipData GetClipData(GPUAnimatorInstancerRuntimeData runtimeData, AnimationClip animationClip)
        {
            if (animationClip == null)
            {
                return default;
            }
            runtimeData.animationClipDataDict.TryGetValue(animationClip.GetHashCode(), out GPUAnimationClipData clipData);
            return clipData;
        }
        
        public float GetClipTime(GPUAnimatorInstancerRuntimeData runtimeData, AnimationClip animationClip)
        {
            int animationKey = animationClip.GetHashCode();
            if (!runtimeData.animationClipDataDict.ContainsKey(animationKey))
            {
                Debug.LogError($"{animationClip.name} 没有烘焙，烘焙试试");
                return 0;
            }
            GPUAnimationClipData clipData = runtimeData.animationClipDataDict[animationKey];

            return GetClipTime(clipData);
        }
        
        public float GetClipTime(GPUAnimationClipData clipData)
        {
            if (clipData.clipFrameCount == 0)
            {
                return 0;
            }
            int clipIndex = GetCurrentClipIndex(clipData);
            if (clipIndex >= 0)
            {
                return GetClipTotalTime(clipIndex, Time.time) % clipData.length;
            }
            return 0;
        }
        
        public float GetClipTotalTime(int clipIndex, float currentTime)
        {
            return (currentTime - currentClipStartTimes[clipIndex]) * currentClipSpeeds[clipIndex];
        }

        public float GetClipFrame(int clipIndex, float currentTime, float clipLength, int clipFrameCount, bool isLoopDisabled)
        {
            float clipTotalTime = GetClipTotalTime(clipIndex, currentTime);
            if (isLoopDisabled && clipTotalTime > clipLength)
            {
                return clipFrameCount - 2;
            }
            else
            {
                return ((clipTotalTime / clipLength) % 1.0f) * (clipFrameCount - 1);
            }
        }
        
        public void RemoveFromTransitioningAnimatorsSwapBack(GPUAnimatorInstancerRuntimeData runtimeData)
        {
            if (runtimeData?.transitioningAnimators == null)
            {
                return;
            }
            
            int lastIndex = runtimeData.transitioningAnimators.Count - 1;
            if (lastIndex != transitionIndex && lastIndex >= 0 && transitionIndex >= 0 && transitionIndex < runtimeData.transitioningAnimators.Count)
            {
                GPUAnimator lastAnimator = runtimeData.transitioningAnimators[lastIndex];
                if (lastAnimator == null)
                {
                    return;
                }
                lastAnimator.transitionIndex = transitionIndex;
                runtimeData.transitioningAnimators[transitionIndex] = lastAnimator;
            }

            transitionIndex = -1;
            if (lastIndex >= 0)
            {
                runtimeData.transitioningAnimators.RemoveAt(lastIndex);
            }
        }
        
        
        //获取挂点信息
        // public AnchorInfo GetAnchorInfoByBoneIndex(GPUAnimatorInstancerRuntimeData runtimeData, int boneIndex, int index)
        // {
        //     if (anchorInfos != null)
        //     {
        //         var animationData = runtimeData.animationData[index * 2];
        //         var controllerData = runtimeData.animatorControllerData[index * 4];
        //         int frameIndex = ((int)animationData.x + (int)controllerData.x) * runtimeData.totalBoneCount + boneIndex;
        //         if (anchorInfos.Length > frameIndex)
        //         {
        //             return anchorInfos[frameIndex];
        //         }
        //     }
        //     return default;
        // }
        //
        // //获取挂点信息
        // public AnchorInfo GetAnchorInfoByAnchorNameHash(GPUAnimatorInstancerRuntimeData runtimeData, int nameHash, int index)
        // {
        //     if (anchorNameToBoneIndex != null && anchorNameToBoneIndex.TryGetValue(nameHash, out int boneIndex))
        //     {
        //         return GetAnchorInfoByBoneIndex(runtimeData, boneIndex, index);
        //     }
        //
        //     return default;
        // }
        
        private int GetCurrentClipIndex(GPUAnimationClipData clipData)
        {
            for (int i = 0; i < currentAnimationClipData.Length; i++)
            {
                if (currentAnimationClipData[i].clipFrameCount > 0 && currentAnimationClipData[i].clipIndex == clipData.clipIndex)
                {
                    return i;
                }
            }
            return -1;
        }
        
        private float GetSpeedRelativeStartTime(float currentTime, float previousStartTime, float previousSpeed, float clipLength, float newSpeed)
        {
            return currentTime - (((currentTime - previousStartTime) * previousSpeed) % clipLength) / newSpeed;
        }
    }
}
